#!/usr/bin/env python3
"""
2_calculate_technical.py
技术指标计算程序
"""

import pandas as pd
import numpy as np
import sys
import os

def calculate_sma(data, window):
    """计算简单移动平均"""
    return data.rolling(window=window).mean()

def calculate_ema(data, span):
    """计算指数移动平均"""
    return data.ewm(span=span, adjust=False).mean()

def calculate_rsi(prices, window=14):
    """计算RSI指标"""
    delta = prices.diff()
    gain = (delta.where(delta > 0, 0)).rolling(window=window).mean()
    loss = (-delta.where(delta < 0, 0)).rolling(window=window).mean()
    rs = gain / loss
    rsi = 100 - (100 / (1 + rs))
    return rsi

def calculate_macd(prices, fast=12, slow=26, signal=9):
    """计算MACD指标"""
    ema_fast = calculate_ema(prices, fast)
    ema_slow = calculate_ema(prices, slow)
    macd = ema_fast - ema_slow
    signal_line = calculate_ema(macd, signal)
    histogram = macd - signal_line
    return macd, signal_line, histogram

def calculate_bollinger_bands(prices, window=20, num_std=2):
    """计算布林带"""
    sma = calculate_sma(prices, window)
    std = prices.rolling(window=window).std()
    upper_band = sma + (std * num_std)
    lower_band = sma - (std * num_std)
    return upper_band, lower_band

def calculate_atr(high, low, close, window=14):
    """计算ATR指标"""
    prev_close = close.shift(1)
    tr1 = high - low
    tr2 = abs(high - prev_close)
    tr3 = abs(low - prev_close)
    true_range = pd.DataFrame({'tr1': tr1, 'tr2': tr2, 'tr3': tr3}).max(axis=1)
    atr = true_range.rolling(window=window).mean()
    return atr

def calculate_kdj(high, low, close, window=9):
    """计算KDJ指标"""
    low_min = low.rolling(window=window).min()
    high_max = high.rolling(window=window).max()
    
    rsv = (close - low_min) / (high_max - low_min) * 100
    k = rsv.ewm(alpha=1/3).mean()
    d = k.ewm(alpha=1/3).mean()
    j = 3 * k - 2 * d
    
    return k, d, j

def calculate_obv(close, volume):
    """计算OBV指标"""
    obv = np.where(close > close.shift(1), volume, 
                   np.where(close < close.shift(1), -volume, 0)).cumsum()
    return pd.Series(obv, index=close.index)

def calculate_adx(high, low, close, window=14):
    """计算ADX指标"""
    prev_close = close.shift(1)
    tr = pd.DataFrame({
        'hl': high - low,
        'hc': abs(high - prev_close),
        'lc': abs(low - prev_close)
    }).max(axis=1)
    
    atr = tr.rolling(window=window).mean()
    
    up_move = high - high.shift(1)
    down_move = low.shift(1) - low
    
    plus_dm = np.where((up_move > down_move) & (up_move > 0), up_move, 0)
    minus_dm = np.where((down_move > up_move) & (down_move > 0), down_move, 0)
    
    plus_di = 100 * (pd.Series(plus_dm).rolling(window=window).mean() / atr)
    minus_di = 100 * (pd.Series(minus_dm).rolling(window=window).mean() / atr)
    
    dx = 100 * abs(plus_di - minus_di) / (plus_di + minus_di)
    adx = dx.rolling(window=window).mean()
    
    return adx

def calculate_technical_indicators():
    """计算所有技术指标"""
    
    # 读取输入数据
    input_file = '../csv_output/level1_raw_data.csv'
    if not os.path.exists(input_file):
        print("❌ level1_raw_data.csv 不存在", file=sys.stderr)
        sys.exit(1)
    
    df = pd.read_csv(input_file)
    df['date'] = pd.to_datetime(df['date'])
    df = df.sort_values('date')
    
    # 计算技术指标
    close = df['close']
    high = df['high']
    low = df['low']
    volume = df['volume']
    
    # 均线系统
    df['ma5'] = calculate_sma(close, 5)
    df['ma10'] = calculate_sma(close, 10)
    df['ma20'] = calculate_sma(close, 20)
    df['ma60'] = calculate_sma(close, 60)
    df['ma120'] = calculate_sma(close, 120)
    df['ma250'] = calculate_sma(close, 250)
    
    # 动量指标
    df['rsi14'] = calculate_rsi(close, 14)
    df['macd'], df['macd_signal'], df['macd_histogram'] = calculate_macd(close)
    df['kdj_k'], df['kdj_d'], df['kdj_j'] = calculate_kdj(high, low, close)
    
    # 趋势指标
    df['boll_upper'], df['boll_lower'] = calculate_bollinger_bands(close)
    df['adx14'] = calculate_adx(high, low, close)
    
    # 波动指标
    df['atr14'] = calculate_atr(high, low, close)
    df['volatility'] = close.pct_change().rolling(20).std() * np.sqrt(252)
    
    # 量价指标
    df['obv'] = calculate_obv(close, volume)
    df['volume_ma5'] = calculate_sma(volume, 5)
    df['volume_ma20'] = calculate_sma(volume, 20)
    
    # 信号指标
    df['golden_cross'] = ((df['ma5'] > df['ma20']) & (df['ma5'].shift(1) <= df['ma20'].shift(1))).astype(int)
    df['death_cross'] = ((df['ma5'] < df['ma20']) & (df['ma5'].shift(1) >= df['ma20'].shift(1))).astype(int)
    
    # 保存结果
    df.to_csv('../csv_output/level2_technical.csv', index=False)
    print(f"✅ 技术指标计算完成：{len(df)}条记录，{len(df.columns)}个指标")

def main():
    if not os.path.exists('../csv_output/level1_raw_data.csv'):
        print("❌ 请先运行1_extract_data.py", file=sys.stderr)
        sys.exit(1)
    
    calculate_technical_indicators()
    print("✅ level2_technical.csv 已生成")

if __name__ == "__main__":
    main()